extern crate alloc;

use alloc::boxed::Box;
use crate::Vec;
use core::ops::Deref;

use crate::common::constants::ErrorCode;
use crate::traits::context_manager::{AuthContext, EnrollContext, ContextManager};
use crate::{log_e, log_i};
use crate::traits::crypto_engine::CryptoEngineRegistry;

const MAX_CONTEXT_NUM: usize = 50;

pub struct DefaultContextManager {
    auth_contexts: Vec<Box<dyn AuthContext>>,
    enroll_context: Option<Box<dyn EnrollContext>>,
}

impl DefaultContextManager {
    pub fn new() -> Self {
        Self {
            auth_contexts: Vec::with_capacity(MAX_CONTEXT_NUM),
            enroll_context: None,
        }
    }
}

impl ContextManager for DefaultContextManager {
    fn add_auth_context(&mut self, context: Box<dyn AuthContext>) -> Result<(), ErrorCode> {
        if self.get_auth_context(context.context_id()).is_ok() {
            log_e!("AuthContext already exists: {}", context.context_id());
            return Err(ErrorCode::GeneralError);
        }

        if self.auth_contexts.len() >= MAX_CONTEXT_NUM {
            log_i!("Reached max context limit, removing oldest context");
            self.auth_contexts.remove(0);
        }

        self.auth_contexts.push(context);
        Ok(())
    }

    fn get_auth_context<'a>(&'a self, context_id: u64) -> Result<&'a dyn AuthContext, ErrorCode> {
        for context in &self.auth_contexts {
            if context.context_id() == context_id {
                return Ok(context.as_ref());
            }
        }
        log_e!("AuthContext not found: {}", context_id);
        Err(ErrorCode::GeneralError)
    }

    fn get_auth_context_by_schedule_id<'a>(&'a self, schedule_id: u64) -> Result<&'a dyn AuthContext, ErrorCode> {
        for context in &self.auth_contexts {
            if context.schedule_id() == schedule_id {
                return Ok(context.as_ref());
            }
        }
        log_e!("AuthContext not found for schedule_id: {}", schedule_id);
        Err(ErrorCode::GeneralError)
    }

    fn remove_auth_context(&mut self, context_id: u64) -> Result<(), ErrorCode> {
        let original_len = self.auth_contexts.len();
        self.auth_contexts.retain(|c| c.context_id() != context_id);
        
        if self.auth_contexts.len() < original_len {
            Ok(())
        } else {
            log_e!("context_id not found: {}", context_id);
            Err(ErrorCode::GeneralError)
        }
    }

    fn add_enroll_context(&mut self, context: Box<dyn EnrollContext>) -> Result<(), ErrorCode> {
        if self.enroll_context.is_some() {
            log_e!("EnrollSchedule already exists");
            return Err(ErrorCode::GeneralError);
        }

        self.enroll_context = Some(context);
        Ok(())
    }

    fn get_enroll_context<'a>(&'a self) -> Result<&'a dyn EnrollContext, ErrorCode> {
        match &self.enroll_context {
            Some(context) => Ok(context.as_ref()),
            _ => {
                log_e!("EnrollContext not found");
                Err(ErrorCode::GeneralError)
            }
        }
    }

    fn remove_enroll_context(&mut self) -> Result<(), ErrorCode> {
        match &self.enroll_context {
            Some(_) => {
                self.enroll_context = None;
                Ok(())
            }
            None => {
                log_e!("EnrollContext not found");
                Err(ErrorCode::GeneralError)
            }
        }
    }

    fn generate_unique_schedule_id(&self) -> Result<u64, ErrorCode> {
        let mut existing_ids = Vec::new();
        for auth_context in &self.auth_contexts {
            existing_ids.push(auth_context.schedule_id());
        }

        if let Some(enroll_context) = &self.enroll_context {
            existing_ids.push(enroll_context.schedule_id());
        }

        let mut attempts = 0;
        const MAX_ATTEMPTS: usize = 100;
        loop {
            let mut random_bytes = [0u8; 8];
            CryptoEngineRegistry::get().secure_random(&mut random_bytes)
                .map_err(|_| ErrorCode::GeneralError)?;
            let random = u64::from_ne_bytes(random_bytes);
            let exists = existing_ids.iter().any(|item| *item == random);
            if !exists {
                return Ok(random);
            }

            attempts += 1;
            if attempts >= MAX_ATTEMPTS {
                log_e!("random generate error");
                return Err(ErrorCode::GeneralError);
            }
        }
    }

    fn check_context_id_duplicate(&self, context_id: u64) -> Result<bool, ErrorCode> {
        for context in &self.auth_contexts {
            if context.context_id() == context_id {
                return Ok(true);
            }
        }
        Ok(false)
    }
}
